package com.wanger.swc.util.messagePush;

import com.alibaba.fastjson2.JSON;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

@Slf4j
public class SseEmitterUtil {

    //使用 map 对象，便于根据 userId 来获取对应的 SseEmitter，或者放 redis 里面
    private static final Map<String, SseEmitter> sseEmitterMap = new ConcurrentHashMap<>();


    /**
     * 创建用户连接并返回 SseEmitter
     *
     * @param userId 用户ID
     * @return SseEmitter
     */
    public static SseEmitter connect(String userId) {
        //设置超时时间，0表示不过期。默认30秒，超过时间未完成会抛出异常：AsyncRequestTimeoutException
        // 1天 = 24小时 * 60分钟 * 60秒 * 1000毫秒
        SseEmitter sseEmitter = new SseEmitter(24L * 60 * 60 * 1000);
//        SseEmitter sseEmitter =  new SseEmitter(2000L);
        try {
            //注册回调：完成、失败、超时
            sseEmitter.onCompletion(completionCallBack(userId));
            sseEmitter.onError(errorCallBack(userId));
            sseEmitter.onTimeout(timeoutCallBack(userId));
            // 缓存
            sseEmitterMap.put(userId, sseEmitter);
            log.info("创建新的sse连接，当前用户：{}", userId);
        } catch (Exception e) {
            log.error("创建新的sse连接异常，当前用户：{}", userId, e);
        }
        return sseEmitter;
    }

    /**
     * 给指定用户发送信息
     *
     * @param userId 用户id
     * @param data   消息内容
     */
    public static void sendData(String userId, Object data) {
        sendMessage(userId, JSON.toJSONString(data));
    }

    /**
     * 给指定用户发送信息
     *
     * @param userId  用户id
     * @param message 消息内容
     */
    public static void sendMessage(String userId, String message) {
        if (!sseEmitterMap.containsKey(userId)) return;
        try {
            sseEmitterMap.get(userId).send(message);
        } catch (IOException e) {
            log.error("用户[{}]推送异常:{}", userId, e.getMessage());
            removeUser(userId);
        }
    }

    /**
     * 给所有用户发送信息
     *
     * @param data 消息数据
     */
    public static void sendDataAll(Object data) {
        sendMessageAll(JSON.toJSONString(data));
    }

    /**
     * 给所有用户发送信息
     *
     * @param message 消息内容
     */
    public static void sendMessageAll(String message) {
        sseEmitterMap.keySet().forEach(userId -> sendMessage(userId, message));
    }

    /**
     * 获取所有用户
     *
     * @return 用户列表
     */
    public static Set<String> getSseKey() {
        return sseEmitterMap.keySet();
    }

    /**
     * 移除用户连接
     *
     * @param userId 用户ID
     */
    public static void removeUser(String userId) {
        sseEmitterMap.remove(userId);
        log.info("移除用户：{}", userId);
    }

    /**
     * 获取当前连接数量
     *
     * @return 连接数量
     */
    public static int getUserCount() {
        return sseEmitterMap.size();
    }

    /**
     * 移除连接回调
     *
     * @param userId 用户ID
     * @return 移除连接回调
     */
    private static Runnable completionCallBack(String userId) {
        return () -> {
            log.info("结束连接{}", userId);
            removeUser(userId);
        };
    }

    /**
     * 超时回调
     *
     * @param userId 用户ID
     * @return 超时回调
     */
    private static Runnable timeoutCallBack(String userId) {
        return () -> {
            log.info("连接超时{}", userId);
            removeUser(userId);
        };
    }

    /**
     * 错误回调
     *
     * @param userId 用户ID
     * @return 错误回调
     */
    private static Consumer<Throwable> errorCallBack(String userId) {
        return throwable -> {
            log.error("连接异常{}", userId);
            removeUser(userId);
        };
    }

}
